import os
import cv2
from tqdm import tqdm as tqdm
import numpy as np
from lxml import etree

# path=r'/Users/edy/Downloads/YQSL_Fridge_CptYqsl_image_train_v0'
# anno_path=r"/Users/edy/Downloads/YQSL_Fridge_CptYqsl_image_train_v0/Annotations"
# img_path=r"/Users/edy/Downloads/YQSL_Fridge_CptYqsl_image_train_v0/images/train"
# out_crop_dir=os.path.join(os.path.dirname(anno_path),'crop')

path = r"/Users/edy/Data/RC/label_data/xianchang_test/2022-02-11/test"
anno_path = os.path.join(path, 'Annotations')
# anno_path = os.path.join(path, 'pre1')
img_path = os.path.join(path, 'JPEGImages')

out_crop_dir = os.path.join(path, 'crops')

# target_cat=["东鹏加気能量型营养素饮料罐装355毫升",
# "东鹏特饮维生素功能饮料瓶装250毫升",
# "东鹏特饮维生素功能饮料瓶装500毫升",
# "东鹏特饮维生素功能饮料罐装250毫升"]

target_cat=['依仑大白梨汽水500ml',
'依仑大连汽水500ml',
'依仑大连汽水桔子味500ml',
'依仑大连汽水苹果味500ml',
'依仑大连荔枝味碳酸饮料500ml',
'依仑小蜜蜂葡萄果味型碳酸饮料500ml',
'八王寺山楂香梨维生素果汁汽水380ml',
'八王寺荔枝维生素果汁汽水380ml',
'八王寺香橙维生素果汁汽水380ml',
'大窑嘉宾果味汽水520ml',
'大窑嘉宾果味汽水PET520ml',
'大窑橙诺520ml',
'大窑橙诺PET520ml',
'宏宝莱果汁鲜山楂味饮料玻璃瓶200ml',
'宏宝莱果汁鲜桔子味汽水饮料玻璃瓶200ml',
'宏宝莱果汁鲜脐橙味饮料玻璃瓶200ml',
'宏宝莱生榨山水蜜桃果味饮料300ml',
'宏宝莱生榨蓝莓苹果复合果味饮料300ml',
'宏宝莱生榨香橙果肉饮料300ml',
'宏宝莱老汽水550ml',
'宏宝莱花生露植物蛋白饮料445ml',
'宏宝莱花生露植物蛋白饮料玻璃瓶200ml',
'宏宝莱荔枝味汽水饮料玻璃瓶200ml',
'宏宝莱荔枝汽水330ml铝罐',
'旺仔大白梨果味型碳酸饮料310ml',
'珍珍桃味汽水330ml',
'珍珍桔味汽水330ml',
'珍珍荔枝味汽水330ml',
'秋林格瓦斯PE瓶350ml',
'秋林黑格瓦斯PE瓶350ml',
            ]

target_cat=[]
if os.path.exists(out_crop_dir) is False:
    os.makedirs(out_crop_dir)


def recursive_parse_xml_to_dict(xml):
    if not xml:
        return {xml.tag: xml.text}
    result = {}
    for child in xml:
        child_result = recursive_parse_xml_to_dict(child)
        if child.tag != 'object':
            result[child.tag] = child_result[child.tag]
        else:
            if child.tag not in result:
                result[child.tag] = []
            result[child.tag].append(child_result[child.tag])
    return {xml.tag: result}


anno_list = list(filter(lambda x: x.endswith('.xml'), os.listdir(anno_path)))
for i in tqdm(anno_list):
    xmlfile = os.path.join(anno_path, i)
    with open(xmlfile, 'rb') as f:
        xml_str = f.read()
    # img = cv2.imread(os.path.join(img_path, i.replace('.xml', '.jpg')))
    img = cv2.imdecode(np.fromfile(os.path.join(img_path, i.replace('.xml', '.jpg')),dtype=np.uint8),1)
    xml_dict = recursive_parse_xml_to_dict(etree.fromstring(xml_str))
    # index=0
    if 'object' in xml_dict['annotation']:
        for j in range(len(xml_dict['annotation']['object']) - 1, -1, -1):
            info = xml_dict['annotation']['object'][j]
            name = info['name']

            if name not in target_cat and len(target_cat):  # 不属于目标类的目标不会被crop
                continue

            y1, y2, x1, x2 = int(float(info['bndbox']['ymin'])), int(float(info['bndbox']['ymax'])), int(
                float(info['bndbox']['xmin'])), int(float(info['bndbox']['xmax']))

            crop = img[int(float(info['bndbox']['ymin'])):int(float(info['bndbox']['ymax'])),
                   int(float(info['bndbox']['xmin'])):int(float(info['bndbox']['xmax']))]
            out_path = os.path.join(out_crop_dir, name)
            if os.path.exists(out_path) is False:
                os.makedirs(out_path)
            try:
                cv2.imwrite(os.path.join(out_crop_dir, 'temp.jpg'), crop, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
            except:
                print(xmlfile)
                print(name, y1, x1)
                print('-' * 20)
            else:
                os.rename(os.path.join(out_crop_dir, 'temp.jpg'),
                          os.path.join(out_path, i.replace('.xml', '_{}_{}_{}_{}_.jpg'.format(y1, x1, y2, x2))))
